import sys
import numpy as np

# ''' 1. load hidden state log
# '''

# HS_log = np.load("./logs/hiddenstate_trajectories.npy")
# print("shape of HS_log: ", HS_log.shape)

# ''' 2. load physical trajectory log
# '''

# PHYS_log = np.load("./logs/physical_trajectories.npy")
# print("shape of PHYS_log: ", PHYS_log.shape)

# ''' 3. load goal event log
# '''

# goal_event_log = np.load("./logs/goal_records.npy")
# print("shape of goal_event_log: ", goal_event_log.shape)

# mat_corr = np.load("./logs/dim_corr_mat1.npy")
# mat_corr = np.load("./logs/mat_32.npy")
# mat_corr = np.load("./logs/dim_corr_mat_rnd32_3.npy")
mat_corr = np.load("./logs/dim_corr_mat_linearness32.npy")
# mat_corr = np.load("./logs/dim_corr_mat_linearness32_r0.npy")
# mat_corr = np.load("./logs/dim_corr_mat_linearness32_max.npy")

print(mat_corr)

# 将 mat_32 作为一个三维曲面绘制出来
import matplotlib.pyplot as plt

# 创建一个三维坐标系
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

# 创建一个网格
x, y = np.meshgrid(np.arange(mat_corr.shape[0]), np.arange(mat_corr.shape[1]))

# # 绘制z=0平面的网格
# ax.plot_wireframe(x, y, np.zeros_like(mat_corr), color='gray', linestyle='--')

# # 绘制三维曲线图
# for i in range(mat_corr.shape[0]):
#     for j in range(0, mat_corr.shape[1]-1):
#         if i!=j and i!=j+1:
#             z = [mat_corr[i][j], mat_corr[i][j+1]]
#             # ax.plot([i, i], [j, j+1], z, color=plt.cm.viridis(z[0]-0.2))
#             ax.plot([i, i], [j, j+1], z, color=plt.cm.viridis(z[0]-2.5))
#         elif i == j+1 and j+2 < mat_corr.shape[1]:
#             z = [mat_corr[i][j], mat_corr[i][j+2]]
#             ax.plot([i, i], [j, j+2], z, color='r')

# 绘制三维曲线图
for i in range(mat_corr.shape[0]):
    for j in range(0, mat_corr.shape[1]-2):
        z = [mat_corr[i][j], mat_corr[i][j+1]]
        ax.plot([i, i], [j, j+1], z, color=plt.cm.viridis(z[0]-2.5))

# 显示图形
plt.show()